package com.joysuccess.snmptrap.util;

import com.joysuccess.common.utils.HelpUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.snmp4j.CommunityTarget;
import org.snmp4j.PDU;
import org.snmp4j.Snmp;
import org.snmp4j.TransportMapping;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.event.ResponseListener;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.smi.*;
import org.snmp4j.transport.DefaultUdpTransportMapping;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * @author shenchao
 */
public class SnmpUtil {
    private static Log logger = LogFactory.getLog(SnmpUtil.class);

    public static final int DEFAULT_VERSION = SnmpConstants.version2c;
    public static final String DEFAULT_PROTOCOL = "udp";
    public static final int DEFAULT_PORT = 161;
    public static final long DEFAULT_TIMEOUT = 3 * 1000L;
    public static final int DEFAULT_RETRY = 3;
    public static final String COMMUNITY = "public";

    /**
     * 创建对象communityTarget，用于返回target
     *
     * @param ip
     * @param port
     * @return CommunityTarget
     */
    public static CommunityTarget createDefault(String ip, int port) {
        Integer snmpPort;
        if (0 == port || HelpUtils.isEmpty(port)) {
            snmpPort = DEFAULT_PORT;
        } else {
            snmpPort = port;
        }
        Address address = GenericAddress.parse(DEFAULT_PROTOCOL + ":" + ip
                + "/" + snmpPort);
        CommunityTarget target = new CommunityTarget();
        target.setCommunity(new OctetString(COMMUNITY));
        target.setAddress(address);
        target.setVersion(DEFAULT_VERSION);
        // milliseconds
        target.setTimeout(DEFAULT_TIMEOUT);
        target.setRetries(DEFAULT_RETRY);
        return target;
    }

    /**
     * 根据OID，获取单条消息
     */
    public static Map<String, String> snmpGet(String ip, int port, String oid) {

        CommunityTarget target = createDefault(ip, port);
        Snmp snmp = null;
        Map<String, String> result = new HashMap<>();
        try {
            PDU pdu = new PDU();
            pdu.add(new VariableBinding(new OID(oid)));
            DefaultUdpTransportMapping transport = new DefaultUdpTransportMapping();
            snmp = new Snmp(transport);
            snmp.listen();
            pdu.setType(PDU.GET);
            ResponseEvent respEvent = snmp.send(pdu, target);
            PDU response = respEvent.getResponse();

            if (HelpUtils.isNotEmpty(response)) {
                for (int i = 0; i < response.size(); i++) {
                    VariableBinding vb = response.get(i);
                    result.put(vb.getOid().toString(), vb.getVariable().toString());
                }

            }
            return result;
        } catch (Exception e) {
            logger.error("SNMP Get Exception:" + e);
            return result;
        } finally {
            if (snmp != null) {
                try {
                    snmp.close();
                } catch (IOException ex1) {
                    snmp = null;
                }
            }
        }
    }

    /**
     * 根据OID列表，一次获取多条OID数据，并且以List形式返回
     */
    public static Map<String, String> snmpGetList(String ip, int port, List<String> oidList) {
        CommunityTarget target = createDefault(ip, port);
        Snmp snmp = null;
        Map<String, String> result = new HashMap<>();
        try {
            PDU pdu = new PDU();

            for (String oid : oidList) {
                pdu.add(new VariableBinding(new OID(oid)));
            }

            DefaultUdpTransportMapping transport = new DefaultUdpTransportMapping();
            snmp = new Snmp(transport);
            snmp.listen();
            pdu.setType(PDU.GET);
            ResponseEvent respEvent = snmp.send(pdu, target);
            PDU response = respEvent.getResponse();

            if (HelpUtils.isNotEmpty(response)) {
                getResultFromResponse(response, result);
            }
            return result;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("SNMP GetList Exception:" + e);
            return result;

        } finally {
            if (snmp != null) {
                try {
                    snmp.close();
                } catch (IOException ex1) {
                    snmp = null;
                }
            }
        }
    }

    /**
     * 根据OID列表，采用异步方式一次获取多条OID数据，并且以List形式返回
     */
    public static Map<String, String> snmpAsynGetList(String ip, int port, List<String> oidList) {
        CommunityTarget target = createDefault(ip, port);
        Snmp snmp = null;
        Map<String, String> result = new HashMap<>();
        try {
            PDU pdu = new PDU();

            for (String oid : oidList) {
                pdu.add(new VariableBinding(new OID(oid)));
            }

            DefaultUdpTransportMapping transport = new DefaultUdpTransportMapping();
            snmp = new Snmp(transport);
            snmp.listen();
            pdu.setType(PDU.GET);
            ResponseEvent respEvent = snmp.send(pdu, target);
            PDU response = respEvent.getResponse();

            /*异步获取*/
            final CountDownLatch latch = new CountDownLatch(1);
            ResponseListener listener = new ResponseListener() {
                @Override
                public void onResponse(ResponseEvent event) {
                    ((Snmp) event.getSource()).cancel(event.getRequest(), this);
                    PDU response = event.getResponse();
                    PDU request = event.getRequest();
                    if (HelpUtils.isNotEmpty(response) && response.getErrorStatus() == 0) {
                        getResultFromResponse(response, result);
                        latch.countDown();
                    }
                }
            };

            pdu.setType(PDU.GET);
            snmp.send(pdu, target, null, listener);
            boolean wait = latch.await(30, TimeUnit.SECONDS);
            snmp.close();
            return result;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("SNMP AsynGetList Exception:" + e);
            return result;
        } finally {
            if (snmp != null) {
                try {
                    snmp.close();
                } catch (IOException ex1) {
                    snmp = null;
                }
            }

        }
    }

    private static void getResultFromResponse(PDU response, Map<String, String> result) {
        for (int i = 0; i < response.size(); i++) {
            VariableBinding vb = response.get(i);
            result.put(vb.getOid().toString(), vb.getVariable().toString());
        }
    }

    /**
     * 根据targetOID，获取树形数据
     */
    public static Map<String, String> snmpWalk(String ip, int port, String targetOid) {
        CommunityTarget target = createDefault(ip, port);
        TransportMapping transport = null;
        Snmp snmp = null;
        Map<String, String> result = new HashMap<>();
        try {
            transport = new DefaultUdpTransportMapping();
            snmp = new Snmp(transport);
            transport.listen();

            PDU pdu = new PDU();
            OID targetOID = new OID(targetOid);
            pdu.add(new VariableBinding(targetOID));

            boolean finished = false;
            while (!finished) {
                VariableBinding vb = null;
                ResponseEvent respEvent = snmp.getNext(pdu, target);

                PDU response = respEvent.getResponse();

                if (HelpUtils.isEmpty(response)) {
                    finished = true;
                    break;
                } else {
                    vb = response.get(0);
                }
                // check finish
                finished = checkWalkFinished(targetOID, pdu, vb);
                if (!finished) {
                    result.put(vb.getOid().toString(), vb.getVariable().toString());
                    // Set up the variable binding for the next entry.
                    pdu.setRequestID(new Integer32(0));
                    pdu.set(0, vb);
                } else {
                    snmp.close();
                }
            }
            return result;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("SNMP walk Exception: " + e);
            return result;
        } finally {
            if (snmp != null) {
                try {
                    snmp.close();
                } catch (IOException ex1) {
                    snmp = null;
                }
            }
        }
    }

    private static boolean checkWalkFinished(OID targetOID, PDU pdu,
                                             VariableBinding vb) {
        boolean finished = false;
        if (pdu.getErrorStatus() != 0) {
            finished = true;
        } else if (vb.getOid() == null) {
            finished = true;
        } else if (vb.getOid().size() < targetOID.size()) {
            finished = true;
        } else if (targetOID.leftMostCompare(targetOID.size(), vb.getOid()) != 0) {
            finished = true;
        } else if (Null.isExceptionSyntax(vb.getVariable().getSyntax())) {
            finished = true;
        } else if (vb.getOid().compareTo(targetOID) <= 0) {
            finished = true;
        }
        return finished;

    }

    /**
     * 根据targetOID，异步获取树形数据
     */
    public static Map<String, String> snmpAsynWalk(String ip, int port, String oid) {
        final CommunityTarget target = createDefault(ip, port);
        Snmp snmp = null;
        Map<String, String> result = new HashMap<>();
        try {
            DefaultUdpTransportMapping transport = new DefaultUdpTransportMapping();
            snmp = new Snmp(transport);
            snmp.listen();

            final PDU pdu = new PDU();
            final OID targetOID = new OID(oid);
            final CountDownLatch latch = new CountDownLatch(1);
            pdu.add(new VariableBinding(targetOID));

            ResponseListener listener = new ResponseListener() {
                @Override
                public void onResponse(ResponseEvent event) {
                    ((Snmp) event.getSource()).cancel(event.getRequest(), this);

                    try {
                        PDU response = event.getResponse();
                        if (HelpUtils.isNotEmpty(response) && response.getErrorStatus() == 0) {
                            VariableBinding vb = response.get(0);

                            boolean finished = checkWalkFinished(targetOID, pdu, vb);
                            if (!finished) {
                                result.put(vb.getOid().toString(), vb.getVariable().toString());
                                pdu.setRequestID(new Integer32(0));
                                pdu.set(0, vb);
                                ((Snmp) event.getSource()).getNext(pdu, target,
                                        null, this);
                            } else {
                                latch.countDown();
                            }

                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                        logger.error("SNMP Asyn Walk Exception:" + e);
                        latch.countDown();
                    }

                }
            };
            snmp.getNext(pdu, target, null, listener);
            boolean wait = latch.await(30, TimeUnit.SECONDS);
            snmp.close();
            return result;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("SNMP Asyn Walk Exception:" + e);
            return result;
        }
    }

    /**
     * 根据OID和指定string来设置设备的数据
     */
    public static void setPDU(String ip, String oid, int port, String val) throws IOException {
        CommunityTarget target = createDefault(ip, port);
        Snmp snmp = null;
        PDU pdu = new PDU();
        pdu.add(new VariableBinding(new OID(oid), new OctetString(val)));
        pdu.setType(PDU.SET);

        DefaultUdpTransportMapping transport = new DefaultUdpTransportMapping();
        snmp = new Snmp(transport);
        snmp.listen();
        snmp.send(pdu, target);
        snmp.close();
    }

    /**
     * 解析OctetString
     * @param octetString
     * @return
     */
    public static String dealOctetStr(String octetString, String charset){
        try{
            if(!octetString.contains(":")) {
                return octetString;
            }
            String[] temps = octetString.split(":");
            byte[] bs = new byte[temps.length];
            for(int i=0;i<temps.length;i++)
                bs[i] = (byte) Integer.parseInt(temps[i],16);
            return new String(bs,charset);
        }catch(Exception e){
            return null;
        }
    }

    public static void main(String[] args) {
        Map<String, String> map = new HashMap<>();
//        map = snmpAsynGetList("10.1.0.2",161, Arrays.asList("1.3.6.1.4.1.93450.1.122.51","1.3.6.1.4.1.93450.1.122.50"));
//        map = snmpGet("10.1.0.2",161,"1.3.6.1.4.1.93450.1.122.51");
        System.out.println(map.size());
    }
}